import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import random

# 设置随机种子确保结果可复现
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

# 检查是否有可用的GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# 数据预处理：将图像转换为张量并标准化
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为Tensor，范围从[0,255]变为[0,1]
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差
])

# 加载MNIST数据集
train_dataset = torchvision.datasets.MNIST(
    root='./data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(
    root='./data', train=False, download=True, transform=transform)

# 创建数据加载器
batch_size = 64
train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False)


# 定义两层神经网络模型
class TwoLayerNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(TwoLayerNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)  # 第一个全连接层
        self.relu = nn.ReLU()  # 激活函数
        self.fc2 = nn.Linear(hidden_size, num_classes)  # 第二个全连接层

    def forward(self, x):
        # 将输入张量展平为一维向量
        x = x.view(-1, input_size)
        # 前向传播
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out


# 模型参数
input_size = 784  # 28x28
hidden_size = 128  # 隐藏层神经元数量
num_classes = 10  # 0-9共10个数字

# 初始化模型并移至GPU（如果可用）
model = TwoLayerNet(input_size, hidden_size, num_classes).to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
learning_rate = 0.01
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)

# 训练模型
num_epochs = 10
train_losses = []
train_accuracies = []
test_losses = []
test_accuracies = []

for epoch in range(num_epochs):
    # 训练阶段
    model.train()
    train_loss = 0.0
    correct = 0
    total = 0

    for i, (images, labels) in enumerate(train_loader):
        # 将数据移至GPU（如果可用）
        images, labels = images.to(device), labels.to(device)

        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 统计
        train_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{len(train_loader)}], '
                  f'Loss: {loss.item():.4f}, Accuracy: {100. * correct / total:.2f}%')

    # 计算训练集的平均损失和准确率
    train_loss = train_loss / len(train_loader)
    train_accuracy = 100. * correct / total
    train_losses.append(train_loss)
    train_accuracies.append(train_accuracy)

    # 测试阶段
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()

    # 计算测试集的平均损失和准确率
    test_loss = test_loss / len(test_loader)
    test_accuracy = 100. * correct / total
    test_losses.append(test_loss)
    test_accuracies.append(test_accuracy)

    print(f'Epoch [{epoch + 1}/{num_epochs}], Train Loss: {train_loss:.4f}, '
          f'Train Acc: {train_accuracy:.2f}%, Test Loss: {test_loss:.4f}, '
          f'Test Acc: {test_accuracy:.2f}%')

# 绘制损失曲线
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(train_losses, label='Training Loss')
plt.plot(test_losses, label='Testing Loss')
plt.title('Loss Curves')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()

# 绘制准确率曲线
plt.subplot(1, 2, 2)
plt.plot(train_accuracies, label='Training Accuracy')
plt.plot(test_accuracies, label='Testing Accuracy')
plt.title('Accuracy Curves')
plt.xlabel('Epoch')
plt.ylabel('Accuracy (%)')
plt.legend()
plt.tight_layout()
plt.show()


# 可视化预测结果和真实标签
def visualize_predictions(model, test_loader, num_samples=10):
    model.eval()
    fig = plt.figure(figsize=(15, 5))

    for i in range(num_samples):
        # 随机选择一个样本
        idx = random.randint(0, len(test_dataset) - 1)
        image, true_label = test_dataset[idx]

        # 将图像转换为原始形式（去标准化）
        image_original = image * 0.3081 + 0.1307

        # 预测
        with torch.no_grad():
            image_tensor = image.unsqueeze(0).to(device)
            output = model(image_tensor)
            _, predicted = output.max(1)
            predicted_label = predicted.item()

        # 显示图像和预测结果
        ax = fig.add_subplot(2, 5, i + 1)
        ax.imshow(image_original.squeeze().numpy(), cmap='gray')
        ax.set_title(f'真实: {true_label}, 预测: {predicted_label}',
                     color='green' if true_label == predicted_label else 'red')
        ax.axis('off')

    plt.tight_layout()
    plt.show()


# 可视化10个随机样本的预测结果
visualize_predictions(model, test_loader)

# 输出最终测试准确率
print(f"最终测试集准确率: {test_accuracies[-1]:.2f}%")